'''
https://leetcode.cn/problems/swim-in-rising-water/description/
'''
import heapq
from typing import List


class UnionFind:
    def __init__(self, n):
        self.father = [i for i in range(n)]

    def find(self, i):
        if self.father[i] != i:
            self.father[i] = self.find(self.father[i])
        return self.father[i]

    def is_same_set(self, i, j):
        return self.find(i) == self.find(j)

    def union(self, i, j):
        self.father[self.find(i)] = self.find(j)


class Solution:
    # dijkstra
    def swimInWater(self, grid: List[List[int]]) -> int:
        m, n = len(grid), len(grid[0])
        distance = [[float('inf')] * n for _ in range(m)]
        distance[0][0] = grid[0][0]
        visited = [[False] * n for _ in range(m)]

        heap = [(grid[0][0], 0, 0)]
        while heap:
            cost, x, y = heapq.heappop(heap)
            if not visited[x][y]:
                visited[x][y] = True
                for tx, ty in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]:
                    if 0 <= tx < m and 0 <= ty < n:
                        n_cost = max(cost, grid[tx][ty])
                        if not visited[tx][ty] and n_cost < distance[tx][ty]:
                            distance[tx][ty] = n_cost
                            heapq.heappush(heap, (n_cost, tx, ty))
        return distance[-1][-1]

    # 并查集
    def swimInWater2(self, grid: List[List[int]]) -> int:
        n = len(grid)
        times = n * n
        idx = [(0, 0)] * times       # 索引表，time t and height t 所在网格中的位置
        for i in range(n):
            for j in range(n):
                idx[grid[i][j]] = (i, j)
                
        uf = UnionFind(times)
        for t in range(times):
            x, y = idx[t]
            for nx, ny in [(x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)]:
                if 0 <= nx < n and 0 <= ny < n:
                    if grid[nx][ny] <= t:
                        uf.union(x * n + y, nx * n + ny)
            if uf.is_same_set(0, times - 1):
                return t
        return -1
